#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <stack>
#include <map>
#include <set>
#include <stdexcept>
#include <string>
enum TokenType {
    KW,            // Keyword
    IDN,           // Identifier
    OP,            // Operator
    SE,            // Separator
    INT_LITERAL,   // Integer literal
    FLOAT_LITERAL, // Floating-point literal
    END            // End of input marker
};

struct Token {
    std::string lexeme;    // The exact string from source
    TokenType    type;     // Token category
    std::string  attribute; // Additional attribute (e.g., keyword value or literal text)
};

using namespace std;

// Global token stream and pointer
static vector<Token> tokens;
static size_t current = 0;
static ofstream outputFile;
static int step = 1;

// First and Follow sets
static map<string, set<string>> firstSets = {
    {"compUnit",   {"const", "int", "float", "main"}},
    {"decl",       {"const", "int", "float"}},
    {"valDecl",    {"const"}},
    {"varDecl",    {"int", "float"}},
    {"btype",      {"int", "float"}},
    {"varDef",     {"IDN"}},
    {"funcDef",    {"int", "float", "main"}},
    {"compound_stmt", {"{"}},
    {"statement",  {"if", "return", "{", "IDN", "INT", "FLOAT", "(", ";", "main"}},
    {"expression", {"(", "IDN", "INT", "FLOAT", "main"}},
    {"term",       {"(", "IDN", "INT", "FLOAT", "main"}},
    {"factor",     {"(", "IDN", "INT", "FLOAT", "main"}}
};

static map<string, set<string>> followSets = {
    {"compUnit",   {"EOF"}},
    {"decl",       {"int", "float", "const", "main", "EOF", "}"}},
    {"varDef",     {",", ";"}},
    {"funcDef",    {"int", "float", "const", "main", "EOF"}},
    {"compound_stmt", {"}"}},
    {"statement",  {"}", ";", "else"}}
};

// Helper functions for set checking
static bool inFirst(const string& nonTerminal, const Token& tok) {
    if (firstSets[nonTerminal].count(tok.lexeme)) return true;

    // 处理类型符号（IDN, INT, FLOAT）
    if (tok.type == IDN && firstSets[nonTerminal].count("IDN")) return true;
    if (tok.type == INT_LITERAL && firstSets[nonTerminal].count("INT")) return true;
    if (tok.type == FLOAT_LITERAL && firstSets[nonTerminal].count("FLOAT")) return true;

    // 特殊处理 varDef 和 funcDef
    if (nonTerminal == "varDef" && tok.type == IDN) return true;
    if (nonTerminal == "funcDef" && tok.lexeme == "main") return true;

    // 处理表达式中的标识符和字面量
    if ((nonTerminal == "expression" || nonTerminal == "term" || nonTerminal == "factor") &&
        (tok.type == IDN || tok.type == INT_LITERAL || tok.type == FLOAT_LITERAL)) {
        return true;
    }

    return false;
}

static bool inFollow(const string& nonTerminal, const Token& tok) {
    return followSets[nonTerminal].count(tok.lexeme) || 
           (tok.type == END && followSets[nonTerminal].count("EOF"));
}

// Output the reduction sequence (modified)
void outputReduction(const string& topSymbol, const string& nextToken, const string& action) {
    string nextTokenStr = nextToken.empty() ? "EOF" : nextToken;
    outputFile << step << "\t" << topSymbol << "#" << nextTokenStr << "\t" << action << endl;
    step++;
}

// Helpers
static Token peek() {
    if (current >= tokens.size()) return Token{"", END, ""};
    return tokens[current];
}

static Token advanceTok() {
    if (current >= tokens.size()) return Token{"", END, ""};
    return tokens[current++];
}

static bool match(const vector<string>& opts) {
    for (auto& s : opts) {
        if (peek().lexeme == s) { 
            outputReduction(s, peek().lexeme, "move");
            advanceTok(); 
            return true; 
        }
    }
    return false;
}

static void error(const string& msg) {
    outputReduction("ERROR", peek().lexeme, "error");
    ostringstream oss;
    oss << "Syntax error at token '" << peek().lexeme << "': " << msg;
    throw runtime_error(oss.str());
}

// Forward declarations
void compUnit();
void declaration();
void funcDef();
void compound_stmt();
void statement();
void funcDefTail();
void varDefTail();
void externalDef();
string expression();
string term();
string factor();
bool isNextFunctionDef();

// 新增辅助函数：判断是变量声明还是函数定义
bool isNextFunctionDef() {
    // 保存当前位置
    size_t savedPos = current;
    
    // 跳过类型
    if (tokens[current].lexeme == "int" || tokens[current].lexeme == "float") {
        current++;
    }
    
    // 检查是否是main函数，这是一种特殊情况
    if (current < tokens.size() && tokens[current].lexeme == "main") {
        current++;
        // 如果下一个是左括号，则是函数定义
        bool isFunc = (current < tokens.size() && tokens[current].lexeme == "(");
        // 恢复位置
        current = savedPos;
        return isFunc;
    }
    
    // 检查普通标识符
    if (current < tokens.size() && tokens[current].type == IDN) {
        current++;
        
        // 如果下一个是左括号，则是函数定义
        bool isFunc = (current < tokens.size() && tokens[current].lexeme == "(");
        
        // 恢复位置
        current = savedPos;
        return isFunc;
    }
    
    // 恢复位置
    current = savedPos;
    return false;
}

// Entry
bool runParser(const vector<Token>& inputTokens, const string& outputFilename) {
    tokens = inputTokens;
    current = 0;
    step = 1;
    
    outputFile.open(outputFilename);
    if (!outputFile.is_open()) {
        cerr << "Error: Could not open output file " << outputFilename << endl;
        return false;
    }
    
    outputReduction("program", peek().lexeme, "reduction");
    compUnit();
    
    if (peek().type != END) {
        error("Extra tokens after program end");
    } else {
        outputReduction("END", "EOF", "accept");
    }
    
    outputFile.close();
    cout << "Parsing completed successfully. Output written to " << outputFilename << endl;
    return true;
}

// compUnit → ( decl | funcDef )* EOF
void compUnit() {
    outputReduction("compUnit", peek().lexeme, "reduction");
    while (inFirst("decl", peek()) || inFirst("funcDef", peek())) {
        externalDef();
    }
    
}

// 新增 ExternalDef 和 ExternalDefTail
void externalDef() {
    if (peek().lexeme == "const") { // 常量声明走declaration
        declaration();
    } else { // 变量声明或函数定义
        size_t savedPos = current;
        advanceTok(); // 跳过类型
        advanceTok(); // 跳过标识符
        if (peek().lexeme == "(") { // 是函数定义
            current = savedPos;
            funcDef();
        } else { // 是变量声明
            current = savedPos;
            declaration();
        }
    }
}

// 函数定义的后半部分
void funcDefTail() {
    // 解析参数列表
    if (peek().lexeme != ")") {
        do {
            // 处理参数类型
            outputReduction("btype", peek().lexeme, "reduction");
            if (!match({"int", "float"})) error("Expected parameter type");
            
            // 处理参数名
            if (peek().type != IDN) error("Expected parameter name");
            outputReduction("Ident", peek().lexeme, "move");
            advanceTok();
        } while (match({","}));  // 逗号分隔多个参数
    }
    
    // 结束参数列表并处理函数体
    if (!match({")"})) error("Expected ')' after parameters");
    compound_stmt();
}

// 变量定义后半段：可能初始化，多个变量等
void varDefTail() {
    if (match({"="})) {
        outputReduction("initVal", peek().lexeme, "reduction");
        expression();
    }
    while (match({","})) {
        if (peek().type != IDN) error("Expected identifier after ','");
        outputReduction("Ident", peek().lexeme, "move");
        advanceTok();
        if (match({"="})) {
            outputReduction("initVal", peek().lexeme, "reduction");
            expression();
        }
    }
    if (!match({";"})) error("Expected ';' after variable declaration");
}


// declaration → 'const' bType varDefList ';' (移除了可选const，因为valDecl已包含const)
void declaration() {
    if (!inFirst("decl", peek())) error("Unexpected start of declaration");
    
    // 区分常量声明和变量声明
    if (peek().lexeme == "const") {
        outputReduction("valDecl", peek().lexeme, "reduction");
        match({"const"});
    } else {
        outputReduction("varDecl", peek().lexeme, "reduction");
    }
    
    // 处理类型
    outputReduction("btype", peek().lexeme, "reduction");
    if (!match({"int", "float"})) error("Invalid type");
    
    // 处理变量定义列表
    do {
        outputReduction("varDef", peek().lexeme, "reduction");
        if (peek().type != IDN) error("Expected identifier");
        outputReduction("Ident", peek().lexeme, "move");
        advanceTok();
        
        if (match({"="})) {
            outputReduction("initVal", peek().lexeme, "reduction");
            expression();
        }
    } while (match({","}));
    
    if (!match({";"})) error("Expected ';' after declaration");
}

// funcDef → (bType)? Ident '(' ')' compound_stmt
void funcDef() {
    if (!inFirst("funcDef", peek())) error("Unexpected function definition");
    outputReduction("funcDef", peek().lexeme, "reduction");
    
    // 处理返回类型
    if (peek().lexeme == "int" || peek().lexeme == "float") {
        outputReduction("btype", peek().lexeme, "reduction");
        if (!match({"int","float"})) error("Invalid return type");
    }
    
    // 处理函数名
    if (peek().lexeme == "main") {
        outputReduction("Ident", "main", "move");  // 明确处理main关键字
        advanceTok();
    } else if (peek().type == IDN) {
        outputReduction("Ident", peek().lexeme, "move");
        advanceTok();
    } else {
        error("Expected function name");
    }
    
    // 关键修改：解析参数列表和函数体
    if (!match({"("})) error("Expected '(' after function name");
    funcDefTail();  // 调用参数列表解析函数
}

// compound_stmt → '{' { declaration | statement } '}'
void compound_stmt() {
    if (!inFirst("compound_stmt", peek())) error("Expected '{'");
    outputReduction("compound_stmt", peek().lexeme, "reduction");
    if (!match({"{"})) error("Missing '{'");
    
    while (peek().lexeme != "}") {
        if (inFirst("decl", peek())) declaration();
        else if (inFirst("statement", peek())) statement();
        else error("Invalid block content");
    }
    
    if (!match({"}"})) error("Expected '}'");
}

// statement → if_stmt | return_stmt | expr_stmt
void statement() {
    outputReduction("statement", peek().lexeme, "reduction");
    
    if (peek().lexeme == "if") {
        outputReduction("if_stmt", peek().lexeme, "reduction");
        if (!match({"if"})) error("Expected 'if'");
        if (!match({"("})) error("Expected '(' after 'if'");
        expression();
        if (!match({")"})) error("Expected ')' after condition");
        statement();
        
        // 处理 else 分支
        if (peek().lexeme == "else") {
            outputReduction("else_stmt", peek().lexeme, "reduction");
            if (!match({"else"})) error("Expected 'else'");
            statement();
        }
    }
    else if (peek().lexeme == "return") {
        outputReduction("return_stmt", peek().lexeme, "reduction");
        if (!match({"return"})) error("Expected 'return'");
        
        // 返回语句可能是空的，如：return;
        if (peek().lexeme != ";") {
            expression();
        }
        
        if (!match({";"})) error("Expected ';' after return");
    }
    else {
        outputReduction("expr_stmt", peek().lexeme, "reduction");
        
        // 处理空语句，例如单独的分号
        if (peek().lexeme == ";") {
            match({";"}); // 消耗分号
        } else {
            // 先检查是否是赋值语句
            if (peek().type == IDN) {
                string identName = peek().lexeme;
                outputReduction("Ident", identName, "move");
                advanceTok();
                
                if (match({"="})) {
                    // 这是赋值语句
                    outputReduction("assign", "=", "reduction");
                    expression();
                    if (!match({";"})) error("Expected ';' after assignment");
                    return;
                } else {
                    // 不是赋值语句，回退已消耗的标识符
                    current--;
                }
            }
            
            // 如果不是赋值语句，处理一般表达式
            expression();
            if (!match({";"})) error("Expected ';' after expression");
        }
    }
}

// expression → term ( ('+'|'-'|'=='|'!='|'<'|'>'|'<='|'>=') term )*
string expression() {
    // 处理赋值表达式的情况 (ID = expression)
    if (peek().type == IDN) {
        size_t savedPos = current;
        // 尝试看看是否为赋值表达式
        outputReduction("Ident", peek().lexeme, "move");
        advanceTok();
        
        if (peek().lexeme == "=") {
            // 如果是赋值表达式，往前进一步解析
            outputReduction("=", peek().lexeme, "move");
            advanceTok();
            expression();
            return "int"; // 简化处理，默认为int类型
        } else {
            // 不是赋值表达式，回退
            current = savedPos;
        }
    }
    
    outputReduction("expression", peek().lexeme, "reduction");
    if (!inFirst("expression", peek())) error("Invalid expression start");
    
    string ltype = term();
    
    // 注意：这里加入了所有可能的操作符，确保包含 '>' 符号
    while (match({"+", "-", "==", "!=", "<", ">", "<=", ">=", "&&", "||"})) {
        string op = tokens[current-1].lexeme;
        outputReduction("operator", op, "reduction");
        string rtype = term();
        // 这里可以添加类型检查逻辑
    }
    return ltype;
}

// term → factor ( ('*'|'/'|'&&'|'||') factor )*
string term() {
    outputReduction("term", peek().lexeme, "reduction");
    if (!inFirst("term", peek())) error("Invalid term start");
    
    string ltype = factor();
    
    // 这里我们将 '&&' 和 '||' 移到了 expression 中，因为它们是逻辑运算符，优先级低于算术运算符
    while (match({"*", "/", "%"})) {
        string op = tokens[current-1].lexeme;
        outputReduction("operator", op, "reduction");
        string rtype = factor();
        // 这里可以添加类型检查逻辑
    }
    return ltype;
}

// factor → '(' expression ')' | IDN | Literal
string factor() {
    outputReduction("factor", peek().lexeme, "reduction");
    if (!inFirst("factor", peek())) error("Invalid factor start");
    
    if (match({"("})) {
        string t = expression();
        if (!match({")"})) error("Expected ')'");
        return t;
    }
    else if (peek().type == IDN) {
        string identName = peek().lexeme; // 保存标识符名称
        outputReduction("Ident", identName, "move");
        advanceTok();
        return "int"; // 简化处理，默认标识符为int类型
    }
    else if (peek().type == INT_LITERAL) {
        string value = peek().lexeme; // 保存字面量值
        outputReduction("IntConst", value, "move");
        advanceTok();
        return "int";
    }
    else if (peek().type == FLOAT_LITERAL) {
        string value = peek().lexeme; // 保存字面量值
        outputReduction("FloatConst", value, "move");
        advanceTok();
        return "float";
    }
    
    error("Invalid factor: " + peek().lexeme);
    return "";
}

// 辅助函数：从文件读取token
vector<Token> readTokensFromFile(const string& filename) {
    ifstream inFile(filename);
    if (!inFile.is_open()) {
        cerr << "Error: Could not open token file " << filename << endl;
        exit(1);
    }

    vector<Token> tokens;
    string line;
    while (getline(inFile, line)) {
        if (line.empty()) continue;

        // 处理 EOF 或 END
        if (line.find("<EOF>") != string::npos || line.find("<END>") != string::npos) {
            tokens.push_back({"", END, ""});
            break;
        }

        size_t typeStart = line.find('<');
        size_t typeEnd = line.find('>');

        if (typeStart == string::npos || typeEnd == string::npos) {
            cerr << "Error: Malformed token line: " << line << endl;
            exit(1);
        }

        string lexeme = line.substr(0, typeStart);
        lexeme.erase(0, lexeme.find_first_not_of(" \t"));
        lexeme.erase(lexeme.find_last_not_of(" \t") + 1);

        string typeAndLine = line.substr(typeStart + 1, typeEnd - typeStart - 1);
        size_t commaPos = typeAndLine.find(',');

        TokenType type;
        string lineNum;

        if (commaPos != string::npos) {
            string typeStr = typeAndLine.substr(0, commaPos);
            // 去除 typeStr 前后的空格
            typeStr.erase(0, typeStr.find_first_not_of(" \t"));
            typeStr.erase(typeStr.find_last_not_of(" \t") + 1);

            lineNum = typeAndLine.substr(commaPos + 1);
            lineNum.erase(0, lineNum.find_first_not_of(" \t"));
            lineNum.erase(lineNum.find_last_not_of(" \t") + 1);

            if (typeStr == "KW") {
                type = KW;
                if (lexeme == "main") type = KW; // main 是关键字
            } else if (typeStr == "IDN") type = IDN;
            else if (typeStr == "INT") type = INT_LITERAL;
            else if (typeStr == "FLOAT") type = FLOAT_LITERAL;
            else if (typeStr == "OP") type = OP;
            else if (typeStr == "SE") type = SE;
            else if (typeStr == "EOF" || typeStr == "END") type = END;
            else {
                cerr << "Error: Unknown token type: " << typeStr << endl;
                exit(1);
            }
        } else {
            // 处理无逗号的情况（如 EOF）
            string typeStr = typeAndLine;
            typeStr.erase(0, typeStr.find_first_not_of(" \t"));
            typeStr.erase(typeStr.find_last_not_of(" \t") + 1);

            if (typeStr == "EOF" || typeStr == "END") {
                type = END;
            } else {
                cerr << "Error: Malformed type info: " << typeAndLine << endl;
                exit(1);
            }
        }

        tokens.push_back({lexeme, type, lineNum});
    }

    if (tokens.empty() || tokens.back().type != END) {
        tokens.push_back({"", END, ""});
    }

    return tokens;
}

int main() {
    const std::string tokensFile = "tokens.txt";
    const std::string outputFilename = "output.txt";

    // 确保 output.txt 存在，如果不存在则创建一个空文件
    {
        std::ofstream ofs(outputFilename, std::ios::app);
        if (!ofs) {
            std::cerr << "无法创建或打开输出文件: " << outputFilename << std::endl;
            return 1;
        }
    }

    // 读取 tokens.txt
    std::vector<Token> tokens;
    try {
        tokens = readTokensFromFile(tokensFile);
    } catch (const std::exception& e) {
        std::cerr << "读取令牌文件失败: " << e.what() << std::endl;
        return 1;
    }

    // 运行解析器并将结果写入 output.txt
    if (!runParser(tokens, outputFilename)) {
        std::cerr << "解析失败" << std::endl;
        return 1;
    }

    std::cout << "解析完成，结果保存在 " << outputFilename << std::endl;
    return 0;
}